from NMTDemo.Vocabulary.Vocabulary import Vocabulary


class SequenceVocabulary(Vocabulary):
    def __init__(self,token_to_idx=None,unk_token="<UNK>",
                 mask_token="<MASK>",begin_seq_token="<BEGIN>",
                 end_seq_token="<END>"):
        super(SequenceVocabulary,self).__init__(token_to_idx)
        self._mask_token=mask_token
        self._unk_token=unk_token
        self._begin_seq_token=begin_seq_token
        self._end_seq_token=end_seq_token

        self.mask_index=self.add_token(self._mask_token)
        self.unk_index=self.add_token(self._unk_token)
        self.begin_seq_index=self.add_token(self._begin_seq_token)
        self.end_seq_index=self.add_token(self._end_seq_token)

    def to_serializable(self):
        contents=super(SequenceVocabulary,self).to_serializable()
        contents.update({'unk_token':self._unk_token,
                         'mask_token':self._mask_token,
                         'begin_seq_token':self._begin_seq_token,
                         'end_seq_token':self._end_seq_token})
        return contents

    def lookup_token(self,token):
        if self.unk_index>=0:
            return self._token_to_idx.get(token,self.unk_index)
        else:
            return self._token_to_idx[token]
